//
// Created by myrica on 2/26/25.
//

#ifndef MLIRASTVISITOR_H
#define MLIRASTVISITOR_H

#include <tuple>
#include <utility>

#include "ASTVisitor.h"
#include "Scope.h"
#include "AST/AST.h"
#include "llvm/IR/Module.h"
#include "llvm/IR/PassManager.h"
#include "llvm/Passes/PassBuilder.h"
#include "llvm/Target/TargetMachine.h"
#include "llvm/Transforms/Utils/Mem2Reg.h"
#include "mlir/Conversion/AffineToStandard/AffineToStandard.h"
#include "mlir/Conversion/ArithToLLVM/ArithToLLVM.h"
#include "mlir/Conversion/ComplexToLLVM/ComplexToLLVM.h"
#include "mlir/Conversion/ControlFlowToLLVM/ControlFlowToLLVM.h"
#include "mlir/Conversion/FuncToLLVM/ConvertFuncToLLVM.h"
#include "mlir/Conversion/LLVMCommon/ConversionTarget.h"
#include "mlir/Conversion/LLVMCommon/TypeConverter.h"
#include "mlir/Conversion/MathToLLVM/MathToLLVM.h"
#include "mlir/Conversion/MemRefToLLVM/MemRefToLLVM.h"
#include "mlir/Conversion/SCFToControlFlow/SCFToControlFlow.h"
#include "mlir/Conversion/UBToLLVM/UBToLLVM.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/ControlFlow/IR/ControlFlow.h"
#include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h"
#include "mlir/Dialect/DLTI/DLTI.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/Func/IR/FuncOps.h.inc"
#include "mlir/Dialect/Linalg/IR/Linalg.h"
#include "mlir/Dialect/Linalg/Passes.h"
#include "mlir/Dialect/Linalg/Passes.h.inc"
#include "mlir/Dialect/Linalg/TransformOps/DialectExtension.h"
#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/MemRef/Transforms/Passes.h"
#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/Dialect/Vector/TransformOps/VectorTransformOps.h"
#include "mlir/IR/Builders.h"
#include "mlir/IR/BuiltinOps.h"
#include "mlir/IR/MLIRContext.h"
#include "mlir/Pass/Pass.h"
#include "mlir/Pass/PassManager.h"
#include "mlir/Transforms/Mem2Reg.h"
#include "mlir/Transforms/Passes.h.inc"
#include "Pass/IrisPasses.h"


namespace ast_visitor {
    using namespace ast;
    inline mlir::Type i32, f32, ptr;

    class MLIR_ASTVisitor {
        mlir::MLIRContext *ctx;
        mlir::OpBuilder *builder;
        mlir::ModuleOp module;
        llvm::LLVMContext *llvm_context;

        std::string filename;
        llvm::TargetMachine *TM;

        Scope scope;

        bool emit_mlir, emit_llvm_dialect;
        mlir::func::FuncOp curr_func;
        mlir::Value curr_func_ret_value;
        mlir::Block *curr_func_ret_block;
        mlir::scf::WhileOp curr_while;
        vector<std::tuple<mlir::Block *, mlir::Block *, mlir::Block *> > while_info{};

    public:
        llvm::Module *M;

        explicit MLIR_ASTVisitor(std::string filename,
                                 llvm::TargetMachine *TM,
                                 bool emit_mlir = false,
                                 bool emit_llvm_dialect = false
        ) : filename(std::move(filename)), TM(TM), emit_mlir(emit_mlir), emit_llvm_dialect(emit_llvm_dialect) {
            ctx = new mlir::MLIRContext;
            ctx->enableMultithreading(false);
            ctx->loadDialect<mlir::func::FuncDialect, mlir::memref::MemRefDialect, mlir::arith::ArithDialect,
                mlir::LLVM::LLVMDialect,
                mlir::scf::SCFDialect, mlir::cf::ControlFlowDialect, mlir::DLTIDialect, mlir::linalg::LinalgDialect>();

            mlir::DialectRegistry registry;
            registry.insert<mlir::memref::MemRefDialect, mlir::LLVM::LLVMDialect, mlir::func::FuncDialect,
                mlir::cf::ControlFlowDialect, mlir::arith::ArithDialect, mlir::linalg::LinalgDialect>();
            registerConvertMemRefToLLVMInterface(registry);
            registerConvertFuncToLLVMInterface(registry);
            mlir::cf::registerConvertControlFlowToLLVMInterface(registry);
            mlir::arith::registerConvertArithToLLVMInterface(registry);
            mlir::ub::registerConvertUBToLLVMInterface(registry);
            mlir::linalg::registerTransformDialectExtension(registry);
            mlir::registerConvertMathToLLVMInterface(registry);
            mlir::vector::registerTransformDialectExtension(registry);
            mlir::registerConvertComplexToLLVMInterface(registry);
            ctx->appendDialectRegistry(registry);

            mlir::PassRegistration<irispass::MoveAllocaToEntryPass>();
            mlir::PassRegistration<irispass::RemoveUnreachableOperation>();

            builder = new mlir::OpBuilder(ctx);
            module = builder->create<mlir::ModuleOp>(builder->getUnknownLoc());

            builder->setInsertionPointToEnd(module.getBody());

            i32 = builder->getI32Type();
            f32 = builder->getF32Type();
            ptr = mlir::LLVM::LLVMPointerType::get(ctx);

            M = nullptr;
            llvm_context = new llvm::LLVMContext();
            curr_func = nullptr;
            curr_func_ret_value = nullptr;
            curr_while = nullptr;
        }

        ~MLIR_ASTVisitor() {
            delete builder;
            delete ctx;

            delete M;
            delete llvm_context;
        }

        void visit(std::vector<ASTBase *> *decls);

        void pre_define();

        mlir::Value visitASTBase(ASTBase *base);

        mlir::Value visitExpression(Expression *exp);

        mlir::Value visitExpressionPrimary(ExpressionPrimary *exp);

        mlir::Value visitExpressionPostfix(ExpressionPostfix *exp);

        mlir::Value visitExpressionUnary(ExpressionUnary *exp);

        mlir::Value visitExpressionCast(ExpressionCast *exp);

        mlir::Value visitExpressionMultiplicative(ExpressionMultiplicative *exp);

        mlir::Value visitExpressionAdditive(ExpressionAdditive *exp);

        mlir::Value visitExpressionShift(ExpressionShift *exp);

        mlir::Value visitExpressionRelational(ExpressionRelational *exp);

        mlir::Value visitExpressionEquality(ExpressionEquality *exp);

        mlir::Value visitExpressionAnd(ExpressionAnd *exp);

        mlir::Value visitExpressionExclusiveOr(ExpressionExclusiveOr *exp);

        mlir::Value visitExpressionInclusiveOr(ExpressionInclusiveOr *exp);

        mlir::Value visitExpressionLogicalAnd(ExpressionLogicalAnd *exp);

        mlir::Value visitExpressionLogicalOr(ExpressionLogicalOr *exp);

        mlir::Value visitExpressionConditional(ExpressionConditional *exp);

        mlir::Value visitExpressionAssignment(ExpressionAssignment *exp);

        mlir::DenseElementsAttr visitExpressionInitializerListGlobal(ExpressionInitializerList *exp,
                                                                     llvm::ArrayRef<long> dims,
                                                                     mlir::Type type);

        void visitExpressionInitializerListLocal(ExpressionInitializerList *exp,
                                                 llvm::ArrayRef<long> dims,
                                                 mlir::Type type, mlir::Value alloca_op);


        mlir::Value visitStatement(Statement *statement);

        mlir::Value visitStatementExpression(StatementExpression *statement);

        mlir::Value visitStatementReturn(StatementReturn *ret);

        mlir::Value visitStatementCompound(StatementCompound *compound);

        mlir::Value visitStatementIfElse(StatementIfElse *selection);

        mlir::Value visitStatementWhile(StatementWhile *_while);

        mlir::Value visitStatementBreak(StatementBreak *_break);

        mlir::Value visitStatementContinue(StatementContinue *_continue);

        mlir::Value visitDeclaration(Declaration *decl);

        mlir::Value visitDeclarationFunction(DeclarationFunction *f);

        mlir::Value visitDeclarationBuiltinType(DeclarationBuiltinType *decl);

        // helpers
        [[nodiscard]] auto getConstant(int value) const {
            return builder->create<mlir::arith::ConstantIntOp>(builder->getUnknownLoc(), value, i32);
        }

        [[nodiscard]] auto getConstant(int64_t value) const {
            return builder->create<mlir::arith::ConstantIndexOp>(builder->getUnknownLoc(), value);
        }

        [[nodiscard]] auto getConstant(float value) const {
            return builder->create<mlir::arith::ConstantFloatOp>(builder->getUnknownLoc(), llvm::APFloat(value),
                                                                 builder->getF32Type());
        }

        template<typename T>
        [[ nodiscard]] T extractFromConstantValue(mlir::Value value) const {
            if constexpr (std::is_same_v<T, int>) {
                return cast<mlir::IntegerAttr>(value.getDefiningOp<mlir::arith::ConstantIntOp>().getValue()).
                        getInt();
            } else if constexpr (std::is_same_v<T, float>) {
                return cast<mlir::FloatAttr>(value.getDefiningOp<mlir::arith::ConstantFloatOp>().getValue())
                        .getValueAsDouble();
            } else { static_assert(false, "only support int and float in getConstantValue"); }
        }

        size_t _generate_init_list_global(const vector<Expression *> &exps, size_t current_align, size_t current_index,
                                          llvm::ArrayRef<long> dims, vector<int> &values);

        size_t _generate_init_list_local(const vector<Expression *> &exps, size_t current_align, size_t current_index,
                                         size_t current_nested, llvm::ArrayRef<long> dims, mlir::Value alloca_op);
    };

    inline mlir::Type ASTType2MLIRType(const type::ASTType &type) {
        if (std::holds_alternative<type::Int *>(type))
            return i32;
        if (std::holds_alternative<type::Float *>(type))
            return f32;
        if (std::holds_alternative<type::Void *>(type))
            return nullptr;

        llvm_unreachable("no other types.");
    }
}

#endif //MLIRASTVISITOR_H
